import math
import numpy as np
import matplotlib.pyplot as plt
import dezero
from dezero import Variable
from dezero import optimizers
import dezero.functions as F
from dezero.models import MLP
from dezero.datasets import Spiral
from dezero.dataloaders import DataLoader

#设置超参和批次
max_epoch = 300
batch_size = 30
hidden_size = 10
lr = 1.0

#训练数据
train_set = Spiral(train=True)
train_loader = DataLoader(train_set, batch_size)
#测试数据
test_set = Spiral(train=False)
test_loader = DataLoader(train_set, batch_size,shuffle=False)

#创建模型 2层网络，优化器
model = MLP((hidden_size, 3))
optimizer = optimizers.SGD(lr).setup(model)

for epoch in range(max_epoch):
    sum_loss, sum_acc = 0, 0
    for x,t in train_loader: #用于训练小批量数据
        # 计算梯度并更新参数
        y = model(x)
        loss = F.softmax_cross_entropy(y, t)
        acc = F.accuracy(y,t)
        model.cleargrads()
        loss.backward()
        optimizer.update()

        sum_loss += float(loss.data) * len(t)
        sum_acc += float(acc.data) * len(t)

    print(f'epoch: {epoch+1}')
    print('train loss: {:.4f}, accuracy: {:.4f}'.format(
        sum_loss / len(train_set), sum_acc / len(train_set)
    ))

    sum_loss, sum_acc = 0, 0
    with dezero.no_grad():
        for x,t in test_loader:
            y = model(x)
            loss = F.softmax_cross_entropy(y, t)
            acc = F.accuracy(y, t)
            sum_loss += float(loss.data) * len(t)
            sum_acc += float(acc.data) * len(t)
    print('test loss: {:.4f}, accuracy: {:.4f}'.format(
        sum_loss / len(train_set), sum_acc / len(train_set)
    ))




# Plot boundary area the model predict
h = 0.001
# 确认训练集的边界
x = train_set.data
x_min, x_max = x[:, 0].min() - .1, x[:, 0].max() + .1
y_min, y_max = x[:, 1].min() - .1, x[:, 1].max() + .1
# 生成网格数据，xx:所有网格点的x坐标，形状也是网格性nxm。yy同样
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h)) #创建网格点
# xx,yy的扁平化成一串坐标点（密密麻麻的网格点平摊开来）
X = np.c_[xx.ravel(), yy.ravel()]
# 对网格点进行类型预测 注意不需要反向传播
with dezero.no_grad():
    score = model(X)
predict_cls = np.argmax(score.data, axis=1)
# 预测类型后，重新变回网格的样子，因为后面contourf接收网格形式的绘图数据
Z = predict_cls.reshape(xx.shape)
plt.contourf(xx, yy, Z)

# Plot data points of the dataset
N, CLS_NUM = 100, 3
markers = ['o', 'x', '^']
colors = ['orange', 'blue', 'green']
for i in range(len(x)):
    c = train_set.label[i]
    plt.scatter(x[i][0], x[i][1], s=40,  marker=markers[c], c=colors[c])
plt.show()
